package com.gitee.feizns.explore.data_structure.tree.cus.bt;

import com.gitee.feizns.explore.data_structure.tree.cus.Tree;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;

/**
 * @author feizns
 * @since 2019/12/11
 */
public class Btree<K extends Comparable<K>, V> implements Tree<K, V> {

    private BNode<K, V> root;

    private int size;

    @Override
    public V get(Object key) {
        BNode tmp = root;
        while (tmp != null) {
            for (int i = 0; tmp != null && i < tmp.size(); i++) {
                Entry e = tmp.entry(i);
                int com = ((K) key).compareTo(e.getKey());
                if (com < 0)
                    tmp = tmp.node(i);
                else if (com > 0 && i == tmp.size() - 1)
                    tmp = tmp.node(tmp.size());
                else if (com == 0)
                    return e.getValue();
            }
        }
        return null;
    }

    @Override
    public V put(K key, V value) {
        NAndV ret = put(root, key, value);
        root = ret.n;
        return ret.v;
    }

    private NAndV put(BNode n, K k, V v) {
        if (n == null) {
            size++;
            return new NAndV(new BNode(new Entry(k, v)), true);
        }

        V oldV = null;
        int nodeIndex = -1;
        for (int i = 0; i < n.size(); i++) {
            Entry e = n.entry(i);
            int com = k.compareTo(e.getKey());

            if (com < 0)
                nodeIndex = i;
            else if (com > 0 && i == n.size() - 1)
                nodeIndex = n.size();
            else if (com == 0)
                oldV = e.setValue(v);
        }

        if (nodeIndex != -1) {
            BNode node = n.node(nodeIndex);
            NAndV ret = put(node, k, v);

            if (ret.merge)
                n = merge(n, ret.n);
            else
                n.childrenNodes.set(nodeIndex, ret.n);

            oldV = ret.v;

            if (n.size() == 4)
                return new NAndV(split(n), oldV, true);
        }

        return new NAndV(n, oldV);
    }

    private BNode merge(BNode original, BNode newNode) {
        List<Entry> entries = original.entries;
        int nextIndex = entries.size();
        Entry e = newNode.entry(0);
        for (int i = 0; i < entries.size(); i++) {
            if (e.getKey().compareTo(entries.get(i).getKey()) < 0) {
                nextIndex = i;
                break;
            }
        }
        List<Entry> newEntries = new ArrayList<>(original.entries);
        newEntries.add(nextIndex, e);
        List<BNode> newNodes = new ArrayList<>(original.childrenNodes);
        List<BNode> nodes = newNode.nodes();
        if (!nodes.isEmpty()) {
            newNodes.remove(nextIndex);
            newNodes.addAll(nextIndex, nodes);
        }
        return new BNode(newEntries, newNodes);
    }

    private BNode split(BNode original) {
        BNode left = new BNode(original.entry(0));
        left.entries.add(original.entry(1));
        left.childrenNodes.set(0, original.node(0));
        left.childrenNodes.set(1, original.node(1));
        left.childrenNodes.set(2, original.node(2));
        BNode right = new BNode(original.entry(3));
        right.childrenNodes.set(0, original.node(3));
        right.childrenNodes.set(1, original.node(4));
        BNode ret = new BNode(original.entry(2));
        ret.childrenNodes.set(0, left);
        ret.childrenNodes.set(1, right);
        return ret;
    }

    @Override
    public V remove(Object key) {
        RNAndV nAndV = remove(root, null, (K) key);
        root = nAndV.n;
        return nAndV.v;
    }

    @Override
    public void clear() {
        root = null;
        size = 0;
    }

    @Override
    public Set<K> keySet() {
        return entrySet().stream().map(Map.Entry::getKey).collect(Collectors.toSet());
    }

    @Override
    public Collection<V> values() {
        return entrySet().stream().map(Map.Entry::getValue).collect(Collectors.toList());
    }

    @Override
    public Set<Map.Entry<K, V>> entrySet() {
        Set<Map.Entry<K, V>> entries = new TreeSet<>();
        recursiveBeforeEach((k, v) -> entries.add(new Entry(k, v)));
        return entries;
    }

    private RNAndV remove(BNode curt, BNode parent, K k) {
        V v;
        if ( curt != null ) {
            int next = curt.size();
            for (int i = 0; i < curt.size(); i++) {
                Entry e = curt.entry(i);
                int comp = k.compareTo(e.k);
                if ( comp < 0 ) {
                    next = i;
                    break;
                } else if ( comp == 0 ) {
                    if ( curt.isLeaf() ) {
                        if ( curt.size() == 1 ) {
                            v = (V) curt.entry(0).v;

                            if ( parent == null ) {
                                size--;
                                return new RNAndV(null, v);
                            }

                            BNode left;
                            BNode right;
                            int nodeIndex = parent.indexOf(curt);
                            if ((left = leftSibling(parent, nodeIndex)) != null && left.size() > 1) {
                                size--;
                                return new RNAndV(rightRotate(parent, nodeIndex), v);
                            } else if ((right = rightSibling(parent, nodeIndex)) != null && right.size() > 1) {
                                size--;
                                return new RNAndV(leftRotate(parent, nodeIndex), v);
                            } else if ( left != null ) {
                                int leftIndex = parent.indexOf(left);
                                Entry en = parent.entry(leftIndex);
                                BNode bNode = merge(left, new BNode(new Entry(en.k, en.v)));
                                parent.entries.remove(leftIndex);
                                parent.childrenNodes.set(leftIndex, bNode);
                                parent.childrenNodes.remove(leftIndex);
                                size--;
                                if (parent.size() == 0)
                                    return new RNAndV(bNode, v);
                                else
                                    return new RNAndV(curt, v);
                            } else if ( right != null ) {
                                int rightIndex = parent.indexOf(right);
                                Entry en = parent.entry(rightIndex - 1);
                                BNode bNode = merge(right, new BNode(new Entry(en.k, en.v)));
                                parent.entries.remove(rightIndex - 1);
                                parent.childrenNodes.set(rightIndex - 1, bNode);
                                parent.childrenNodes.remove(rightIndex);
                                size--;
                                if ( parent.size() == 0 )
                                    return new RNAndV(bNode, v);
                                else
                                    return new RNAndV(curt, v);
                            }
                        } else {
                            v = (V) curt.remove(k);
                            size--;
                            return new RNAndV(curt, v);
                        }
                    } else {
                        v = (V) curt.entry(i).v;
                        BNode right = curt.node(i + 1);
                        Entry min = min(right);
                        RNAndV nAndV = remove(right, curt, min.k);
                        if ( curt.size() == 0 ) {
                            nAndV.n.entries.set(i + 1, new Entry(min.k, min.v));
                            return new RNAndV(nAndV.n, nAndV.v);
                        }
                        curt.entries.set(i, new Entry(min.k, min.v));
                        curt.childrenNodes.set(i + 1, nAndV.n);
                        return new RNAndV(curt, v);
                    }
                }
            }

            //没找到
            RNAndV rnAndV = remove(curt.node(next), curt, k);
            if ( curt.size() == 0 )
                return new RNAndV(rnAndV.n, rnAndV.v);
            curt.childrenNodes.set(next, rnAndV.n);
            return new RNAndV(curt, rnAndV.v);
        }
        return new RNAndV(null, null);
    }

    @Override
    public void recursiveBeforeEach(BiConsumer<K, V> consumer) {
        recursiveBeforeEach(root, consumer);
    }

    @Override
    public void recursiveMiddleEach(BiConsumer<K, V> consumer) {
        recursiveMiddleEach(root, consumer);
    }

    @Override
    public void recursiveAfterEach(BiConsumer<K, V> consumer) {

    }

    private void recursiveMiddleEach(BNode<K, V> node, BiConsumer<K, V> consumer) {
        if ( node != null ) {
            List<Entry> entries = node.entries.stream().filter(Objects::nonNull).collect(Collectors.toList());
            for (int i = 0; i < entries.size(); i++) {
                recursiveMiddleEach(node.node(i), consumer);
                consumer.accept(entries.get(i).getKey(), entries.get(i).getValue());
            }
            recursiveMiddleEach(node.node(entries.size()), consumer);
        }
    }

    private void recursiveBeforeEach(BNode<K, V> node, BiConsumer<K, V> consumer) {
        if ( node != null ) {
            List<Entry> entries = node.entries.stream().filter(Objects::nonNull).collect(Collectors.toList());
            for (int i = 0; i < entries.size(); i++) {
                consumer.accept(entries.get(i).getKey(), entries.get(i).getValue());
                recursiveBeforeEach(node.node(i), consumer);
            }
            recursiveBeforeEach(node.node(entries.size()), consumer);
        }
    }

    private void recursiveAfterEach(BNode<K, V> node, BiConsumer<K, V> consumer) {
//        if ( node != null ) {
//            List<Entry> entries = node.entries.stream().filter(Objects::nonNull).collect(Collectors.toList());
//            for (int i = 0; i < entries.size(); i++) {
//
//            }
//        }
    }

    @Override
    public void recursiveSequenceTraversal(BiConsumer<K, V> consumer) {

    }

    @Override
    public void beforeEach(BiConsumer<K, V> consumer) {

    }

    @Override
    public void middleEach(BiConsumer<K, V> consumer) {

    }

    @Override
    public void afterEach(BiConsumer<K, V> consumer) {

    }

    @Override
    public void sequenceTraversal(BiConsumer<K, V> consumer) {
        Queue<BNode> queue = new LinkedList<>();
        queue.add(root);
        while ( !queue.isEmpty() ) {
            BNode<K, V> node = queue.remove();
            for (Entry entry : node.entries)
                consumer.accept(entry.k, entry.v);
            queue.addAll(node.nodes());
        }
    }

    @Override
    public int size() {
        return size;
    }

    private class RNAndV {

        BNode n;

        V v;

        public RNAndV(BNode n, V v) {
            this.n = n;
            this.v = v;
        }

    }

    public void printTree() {
        printNode(new LinkedList<>(Arrays.asList(root, PLACEHOLDER)));
    }

    private void printNode(Queue<BNode> queue) {
        Queue<BNode> newQueue = new LinkedList<>();
        StringBuilder ret = new StringBuilder();
        while (!queue.isEmpty()) {
            Btree.BNode b = queue.remove();
            StringBuilder tmp = new StringBuilder();
            if ( b == PLACEHOLDER ) {
                tmp.append(String.format("%10s", "|"));
                ret.append(tmp);
                continue;
            }
            tmp.append("[");
            List<Btree.Entry> entries = b.entries;
            for (Btree.Entry entry : entries) {
                if (entry != null)
                    tmp.append(entry.k + ",");
            }
            tmp.setLength(tmp.length() - 1);
            tmp.append("]");
            ret.append(String.format("%10s", tmp));
            newQueue.addAll(b.nodes());
            newQueue.add(PLACEHOLDER);
        }
        System.out.println(String.format("%50s", ret));
        if (!newQueue.isEmpty())
            printNode(newQueue);
    }

    private final BNode PLACEHOLDER = null;

    private BNode rightRotate(BNode p, int nodeIndex) {
        BNode left = p.node(nodeIndex - 1);
        BNode curt = p.node(nodeIndex);
        Entry parent = p.entry(nodeIndex - 1);
        Entry leftLast = left.entry(left.size() - 1);
        left.entries.remove(leftLast);
        curt.entries.set(0, parent);
        p.entries.set(nodeIndex - 1, leftLast);
        return p.node(nodeIndex);
    }

    private BNode leftRotate(BNode p, int nodeIndex) {
        BNode right = p.node(nodeIndex + 1);
        BNode curt = p.node(nodeIndex);
        Entry parent = p.entry(nodeIndex);
        Entry rightFirst = right.entry(0);
        right.entries.remove(rightFirst);
        curt.entries.set(0, parent);
        p.entries.set(nodeIndex, rightFirst);
        return p.node(nodeIndex);
    }

    private BNode leftSibling(BNode p, int index) {
        return p.node(index - 1);
    }

    private BNode rightSibling(BNode p, int index) {
        return p.node(index + 1);
    }

    private Entry min(BNode<K, V> root) {
        BNode curt = root;
        BNode next;
        while ((next = curt.node(0)) != null)
            curt = next;
        return curt.entry(0);
    }


    public class BNode<K, V> {

        public BNode(Entry entry) {
            entries.add(entry);
        }

        public BNode(List<Entry> entries, List<BNode> childrenNodes) {
            this.entries = entries;
            this.childrenNodes = childrenNodes;
        }

        List<Entry> entries = new ArrayList<>();

        List<BNode> childrenNodes = new ArrayList<>();

        {
            childrenNodes.add(null);
            childrenNodes.add(null);
            childrenNodes.add(null);
            childrenNodes.add(null);
        }

        public int size() {
            return entries.size();
        }

        public boolean isEmpty() {
            return size() == 0;
        }

        public Entry entry(int e) {
            return e >= 0 && e < entries.size() ? entries.get(e) : null;
        }

        public BNode<K, V> node(int x) {
            return x >= 0 && x < childrenNodes.size() ? childrenNodes.get(x) : null;
        }

        public List<BNode> nodes() {
            return childrenNodes.stream().filter(Objects::nonNull).collect(Collectors.toList());
        }

        public boolean isLeaf() {
            return childrenNodes.stream().allMatch(Objects::isNull);
        }

        public int indexOf(BNode node) {
            return childrenNodes.indexOf(node);
        }

        public int eIndexOf(K k) {
            for (int i = 0; i < entries.size(); i++) {
                if (entries.get(i).k.equals(k))
                    return i;
            }
            return -1;
        }

        public V remove(K k) {
            V ret = null;
            int index = -1;

            for (int i = 0; i < entries.size(); i++) {
                Entry item = entries.get(i);
                if ( item.k.equals(k) ) {
                    ret = (V) item.v;
                    index = i;
                }
            }

            if (index != -1)
                entries.remove(index);
            return ret;
        }

    }

    class Entry implements Map.Entry<K, V>, Comparable<Entry> {

        K k;

        V v;

        public Entry(K k, V v) {
            this.k = k;
            this.v = v;
        }

        @Override
        public K getKey() {
            return k;
        }

        @Override
        public V getValue() {
            return v;
        }

        @Override
        public V setValue(V value) {
            V o = v;
            v = value;
            return o;
        }

        @Override
        public int compareTo(Entry o) {
            return k.compareTo(o.k);
        }

    }

    private class NAndV {

        boolean merge;

        BNode n;

        V v;

        public NAndV(Btree.BNode n, V v) {
            this.n = n;
            this.v = v;
        }

        public NAndV(BNode n, boolean merge) {
            this.n = n;
            this.merge = merge;
        }

        public NAndV(BNode n, V v, boolean merge) {
            this.n = n;
            this.v = v;
            this.merge = merge;
        }

    }

}
